from contextlib import asynccontextmanager
from typing import AsyncGenerator

import aiomysql
from aiomysql import Cursor, Connection


class AsyncDBError(Exception):
    pass


class AsyncDataBase:
    def __init__(self):
        self.pool: aiomysql.Pool | None = None

    async def connect(self):
        self.pool = await aiomysql.create_pool(
            host="47.108.174.92",
            port=3306,
            user="bili",
            password="123456",
            db="bilibili",
            minsize=5,
            maxsize=20,
            cursorclass=aiomysql.DictCursor,
            autocommit=True,
        )

    async def disconnect(self):
        if self.pool:
            self.pool.close()
            await self.pool.wait_closed()

    @asynccontextmanager
    async def connection(self) -> AsyncGenerator[aiomysql.Connection, None]:
        """获取一个数据库连接"""
        if not self.pool:
            raise AsyncDBError("Database connection pool is not initialized")

        async with self.pool.acquire() as conn:
            try:
                yield conn
            except aiomysql.Error as e:
                raise AsyncDBError(e)
            except Exception as e:
                await conn.rollback()
                raise

    @asynccontextmanager
    async def transaction(self) -> AsyncGenerator[aiomysql.Connection, None]:
        """获取一个带事务的连接"""
        async with self.connection() as conn:
            try:
                yield conn
                await conn.commit()
            except aiomysql.Error as e:
                raise AsyncDBError(e)
            except Exception as e:
                await conn.rollback()
                raise

    @asynccontextmanager
    async def cursor(self) -> AsyncGenerator[aiomysql.Cursor, None]:
        """获取一个游标"""
        async with self.connection() as conn:
            async with conn.cursor() as cur:
                try:
                    yield cur
                finally:
                    # 确保游标关闭
                    pass


asyncDB = AsyncDataBase()


async def get_async_db_cursor() -> AsyncGenerator[Cursor, None]:
    async with asyncDB.cursor() as cursor:
        yield cursor


async def get_async_db_transaction() -> AsyncGenerator[Connection, None]:
    async with asyncDB.transaction() as conn:
        yield conn
